#include "rclcpp/rclcpp.hpp"
#include "sensor_msgs/msg/imu.hpp"
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>
#include <tf2/LinearMath/Transform.h>
#include <tf2/LinearMath/Quaternion.h>
#include <vector>
#include <numeric>
#include <fcntl.h>
#include <termios.h>
#include <unistd.h>


class ImuDriver : public rclcpp::Node
{
private:
    rclcpp::Publisher<sensor_msgs::msg::Imu>::SharedPtr imu_pub;
    sensor_msgs::msg::Imu imu_msg;
    
    // 校准相关变量
    bool is_calibrating = true;
    int calibration_samples = 200;
    int samples_count = 0;
    
    // 零偏数据
    struct BiasData {
        double ax = 0, ay = 0, az = 0;  // 加速度零偏
        double gx = 0, gy = 0, gz = 0;  // 角速度零偏
        double roll = 0, pitch = 0, yaw = 0;  // 姿态零偏
    } bias;

public:
    ImuDriver() : Node("imu_driver")
    {
        // 1. 声明和获取参数
        this->declare_parameter("port", "/dev/ttyUSB0");
        this->declare_parameter("baud_rate", 115200);
        this->declare_parameter("calibration_samples", 200);
        this->declare_parameter("frame_id", "gyro_link");
        this->declare_parameter("topic_name", "/imu/data_raw");
        std::string port = this->get_parameter("port").as_string();
        int baud_rate = this->get_parameter("baud_rate").as_int();
        calibration_samples = this->get_parameter("calibration_samples").as_int();
        std::string frame_id = this->get_parameter("frame_id").as_string();
        std::string topic_name = this->get_parameter("topic_name").as_string();
        RCLCPP_INFO(this->get_logger(), "port: %s", port.c_str());
        RCLCPP_INFO(this->get_logger(), "baud_rate: %d", baud_rate);
        RCLCPP_INFO(this->get_logger(), "calibration_samples: %d", calibration_samples);
        RCLCPP_INFO(this->get_logger(), "frame_id: %s", frame_id.c_str());
        
        // 2. 初始化串口
        if (open_serial(port, baud_rate) != 0) {
            RCLCPP_ERROR(this->get_logger(), "串口初始化失败: %s", port.c_str());
            throw std::runtime_error("串口初始化失败");
        }
        
        // 3. 创建发布者
        imu_pub = this->create_publisher<sensor_msgs::msg::Imu>(topic_name, 10);
        
        // 4. 初始化IMU消息
        imu_msg.header.frame_id = frame_id;
        
        // 5. 设置协方差矩阵
        std::fill(std::begin(imu_msg.orientation_covariance), 
                 std::end(imu_msg.orientation_covariance), 0.0);
        std::fill(std::begin(imu_msg.angular_velocity_covariance), 
                 std::end(imu_msg.angular_velocity_covariance), 0.0);
        std::fill(std::begin(imu_msg.linear_acceleration_covariance), 
                 std::end(imu_msg.linear_acceleration_covariance), 0.0);
        
        // 对角线设置协方差值
        imu_msg.orientation_covariance[0] = 0.0025;
        imu_msg.orientation_covariance[4] = 0.0025;
        imu_msg.orientation_covariance[8] = 0.0025;
        
        imu_msg.angular_velocity_covariance[0] = 0.0025;
        imu_msg.angular_velocity_covariance[4] = 0.0025;
        imu_msg.angular_velocity_covariance[8] = 0.0025;
        
        imu_msg.linear_acceleration_covariance[0] = 0.0025;
        imu_msg.linear_acceleration_covariance[4] = 0.0025;
        imu_msg.linear_acceleration_covariance[8] = 0.0025;
        
        // 6. 重置校准数据
        bias = BiasData();
        
        RCLCPP_INFO(this->get_logger(), "IMU驱动初始化完成，开始校准...");
        RCLCPP_INFO(this->get_logger(), "请保持设备静止 %d 个采样...", calibration_samples);
    }

    void process_imu_data(short* ax, short* gx, short* sAngle)
    {
        if (is_calibrating) {
            // 累积校准数据
            bias.ax += ax[0] / 32768.0 * 16.0 * 9.80665;
            bias.ay += ax[1] / 32768.0 * 16.0 * 9.80665;
            bias.az += (ax[2] / 32768.0 * 16.0 * 9.80665 - 9.80665); // 减去重力加速度

            bias.gx += gx[0] / 32768.0 * 2000 * M_PI / 180.0;
            bias.gy += gx[1] / 32768.0 * 2000 * M_PI / 180.0;
            bias.gz += gx[2] / 32768.0 * 2000 * M_PI / 180.0;

            bias.roll += sAngle[0] / 32768.0 * 180.0 * M_PI / 180.0;
            bias.pitch += sAngle[1] / 32768.0 * 180.0 * M_PI / 180.0;
            bias.yaw += sAngle[2] / 32768.0 * 180.0 * M_PI / 180.0;

            samples_count++;

            if (samples_count >= calibration_samples) {
                // 计算平均值作为零偏
                bias.ax /= calibration_samples;
                bias.ay /= calibration_samples;
                bias.az /= calibration_samples;
                bias.gx /= calibration_samples;
                bias.gy /= calibration_samples;
                bias.gz /= calibration_samples;
                bias.roll /= calibration_samples;
                bias.pitch /= calibration_samples;
                bias.yaw /= calibration_samples;

                is_calibrating = false;
                RCLCPP_INFO(this->get_logger(), "IMU校准完成!");
                RCLCPP_INFO(this->get_logger(), "加速度零偏: x=%.3f, y=%.3f, z=%.3f", 
                           bias.ax, bias.ay, bias.az);
                RCLCPP_INFO(this->get_logger(), "角速度零偏: x=%.3f, y=%.3f, z=%.3f", 
                           bias.gx, bias.gy, bias.gz);
                RCLCPP_INFO(this->get_logger(), "姿态零偏: r=%.3f, p=%.3f, y=%.3f", 
                           bias.roll, bias.pitch, bias.yaw);
            }
            return;
        }

        // 应用零偏校准
        imu_msg.header.stamp = this->now();

        // 补偿加速度
        imu_msg.linear_acceleration.x = ax[0] / 32768.0 * 16.0 * 9.80665 - bias.ax;
        imu_msg.linear_acceleration.y = ax[1] / 32768.0 * 16.0 * 9.80665 - bias.ay;
        imu_msg.linear_acceleration.z = ax[2] / 32768.0 * 16.0 * 9.80665 - bias.az;

        // 补偿角速度
        imu_msg.angular_velocity.x = gx[0] / 32768.0 * 2000 * M_PI / 180.0 - bias.gx;
        imu_msg.angular_velocity.y = gx[1] / 32768.0 * 2000 * M_PI / 180.0 - bias.gy;
        imu_msg.angular_velocity.z = gx[2] / 32768.0 * 2000 * M_PI / 180.0 - bias.gz;

        // 补偿姿态角并转换为四元数
        double roll = sAngle[0] / 32768.0 * 180.0 * M_PI / 180.0 - bias.roll;
        double pitch = sAngle[1] / 32768.0 * 180.0 * M_PI / 180.0 - bias.pitch;
        double yaw = sAngle[2] / 32768.0 * 180.0 * M_PI / 180.0 - bias.yaw;

        tf2::Quaternion quat;
        quat.setRPY(roll, pitch, yaw);
        
        imu_msg.orientation.x = quat.x();
        imu_msg.orientation.y = quat.y();
        imu_msg.orientation.z = quat.z();
        imu_msg.orientation.w = quat.w();

        imu_pub->publish(imu_msg);
    }

    ~ImuDriver()
    {
        close(serial_fd_);
    }

int open_serial(std::string port, int baud_rate)
{
    serial_fd_ = open(port.c_str(), O_RDWR | O_NOCTTY);
    if (serial_fd_ == -1) {
        RCLCPP_ERROR(this->get_logger(), "串口打开失败！");
        return -1;
    }
    
    struct termios options;
    tcgetattr(serial_fd_, &options);
    
    // 设置波特率
    speed_t speed;
    switch(baud_rate) {
        case 9600: speed = B9600; break;
        case 19200: speed = B19200; break;
        case 38400: speed = B38400; break;
        case 57600: speed = B57600; break;
        case 115200: speed = B115200; break;
        default: 
            RCLCPP_ERROR(this->get_logger(), "不支持的波特率！");
            return -1;
    }
    
    // 输入输出波特率
    cfsetispeed(&options, speed);
    cfsetospeed(&options, speed);
    
    // 设置控制模式
    options.c_cflag |= (CLOCAL | CREAD); // 启用接收器并忽略调制解调器控制线
    options.c_cflag &= ~PARENB;          // 无奇偶校验
    options.c_cflag &= ~CSTOPB;          // 1个停止位
    options.c_cflag &= ~CSIZE;           // 清除数据位掩码
    options.c_cflag |= CS8;              // 8位数据
    
    // 设置输入模式
    options.c_iflag &= ~(IXON | IXOFF | IXANY);    // 禁用软件流控制
    options.c_iflag &= ~(INLCR | ICRNL);           // 禁用特殊字符处理
    
    // 设置输出模式
    options.c_oflag &= ~OPOST;                      // 禁用输出处理
    
    // 设置本地模式
    options.c_lflag &= ~(ICANON | ECHO | ECHOE | ISIG); // 原始模式
    
    // 设置等待时间和最小接收字符
    options.c_cc[VTIME] = 1;   // 0.1秒超时
    options.c_cc[VMIN] = 0;    // 无最小字符要求
    
    // 清空缓冲区
    tcflush(serial_fd_, TCIOFLUSH);
    
    // 应用设置
    if (tcsetattr(serial_fd_, TCSANOW, &options) != 0)
    {
        RCLCPP_ERROR(this->get_logger(), "串口设置失败！");
        return -1;
    }
    
    RCLCPP_INFO(this->get_logger(), "串口设置成功！");
    return 0;
}
    int serial_fd_;
};

// main函数修改
int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);
    
    std::shared_ptr<ImuDriver> node;
    try {
        node = std::make_shared<ImuDriver>();
    } catch (const std::exception& e) {
        RCLCPP_ERROR(rclcpp::get_logger("imu_driver"), "初始化失败: %s", e.what());
        return 1;
    }
    
    unsigned char buffer[16];
    short ax[3] = {0};  // 修改为3个元素
    short gx[3] = {0};  // 修改为3个元素
    short sAngle[3] = {0};  // 修改为3个元素

    while (rclcpp::ok())
    {
        // 1. 读取帧头
        int ret = read(node->serial_fd_, buffer, 1);
        if(ret != 1) {
            RCLCPP_WARN(node->get_logger(), "读取帧头失败");
            continue;
        }
        
        // 2. 检查帧头
        if (buffer[0] != 0x55) {
            continue;
        }
        

        read(node->serial_fd_, buffer + 1, 1);
        if (buffer[1] == 0x51)
        {
            read(node->serial_fd_, buffer+2, 1);
            read(node->serial_fd_, buffer+3, 1);
            read(node->serial_fd_, buffer+4, 1);
            read(node->serial_fd_, buffer+5, 1);
            read(node->serial_fd_, buffer+6, 1);
            read(node->serial_fd_, buffer+7, 1);
            read(node->serial_fd_, buffer+8, 1);
            read(node->serial_fd_, buffer+9, 1);
            read(node->serial_fd_, buffer+10, 1);
            unsigned char i, sum;
            for(sum=0,i=0; i<10; i++) sum += buffer[i];
            if(buffer[i] == sum)
            {
                memcpy(ax, buffer+2, 6);
            }
        }
        else if (buffer[1] == 0x52)
        {
            read(node->serial_fd_, buffer+2, 1);
            read(node->serial_fd_, buffer+3, 1);
            read(node->serial_fd_, buffer+4, 1);
            read(node->serial_fd_, buffer+5, 1);
            read(node->serial_fd_, buffer+6, 1);
            read(node->serial_fd_, buffer+7, 1);
            read(node->serial_fd_, buffer+8, 1);
            read(node->serial_fd_, buffer+9, 1);
            read(node->serial_fd_, buffer+10, 1);                
            unsigned char i, sum;
            for(sum=0,i=0; i<10; i++) sum += buffer[i];
            if(buffer[i] == sum)
            {
                memcpy(gx, buffer+2, 6);
            }
        }
        else if (buffer[1] == 0x53)
        {
            read(node->serial_fd_, buffer+2, 1);
            read(node->serial_fd_, buffer+3, 1);
            read(node->serial_fd_, buffer+4, 1);
            read(node->serial_fd_, buffer+5, 1);
            read(node->serial_fd_, buffer+6, 1);
            read(node->serial_fd_, buffer+7, 1);
            read(node->serial_fd_, buffer+8, 1);
            read(node->serial_fd_, buffer+9, 1);
            read(node->serial_fd_, buffer+10, 1);                
            unsigned char i, sum;
            for(sum=0,i=0; i<10; i++) sum += buffer[i];
            if(buffer[i] == sum)
            {
                memcpy(sAngle, buffer+2, 6);
                node->process_imu_data(ax, gx, sAngle);
            }
        }
        
        rclcpp::spin_some(node);
    }
    
    close(node->serial_fd_);
    rclcpp::shutdown();
    return 0;
}
